# @Time : 2021/8/5 8:40
# @Author : Li Kunlun
# @Description : 自定义层


from mxnet import gluon, nd
from mxnet.gluon import nn


# 1、不含模型参数的自定义层
class CenteredLayer(nn.Block):
    def __init__(self, **kwargs):
        super(CenteredLayer, self).__init__(**kwargs)

    def forward(self, x):
        return x - x.mean()


layer = CenteredLayer()
# [-2. -1.  0.  1.  2.]
# <NDArray 5 @cpu(0)>
print(layer(nd.array([1, 2, 3, 4, 5])))

# 用其构造更加复杂的模型
net = nn.Sequential()
net.add(nn.Dense(128),
        CenteredLayer())

# 2、含模型参数的自定义层
params = gluon.ParameterDict()
params.get('param2', shape=(2, 3))
"""
通过get函数从ParameterDict创建Parameter实例
    (
      Parameter param2 (shape=(2, 3), dtype=<class 'numpy.float32'>)
    )
"""
print(params)


class MyDense(nn.Block):
    # units为该层的输出个数，in_units为该层的输入个数
    def __init__(self, units, in_units, **kwargs):
        super(MyDense, self).__init__(**kwargs)
        self.weight = self.params.get('weight', shape=(in_units, units))
        self.bias = self.params.get('bias', shape=(units,))

    def forward(self, x):
        linear = nd.dot(x, self.weight.data()) + self.bias.data()
        return nd.relu(linear)


# 实例化MyDense类并访问它的模型参数
"""
mydense0_ (
  Parameter mydense0_weight (shape=(5, 3), dtype=<class 'numpy.float32'>)
  Parameter mydense0_bias (shape=(3,), dtype=<class 'numpy.float32'>)
)
"""
dense = MyDense(units=3, in_units=5)
print(dense.params)
